
import os, json
from tqdm import tqdm
import random
random.seed(1024)

"""
    测试：交界段之间生成一个过渡句，并给5个过渡句范例。

"""


html_tag_to_extra_id = {
    'title': '<|extra_20|>',
    'h1': '<|extra_21|>',
    'h2': '<|extra_22|>',
    'h3': '<|extra_23|>',
    'h4': '<|extra_24|>',
    'h5': '<|extra_25|>',
    'h6': '<|extra_26|>',
    'strong': '<|extra_27|>',
}


SYSTEM = """
# 角色
你是一名中文系大学生。你可以准确理解题意，并在充分理解全文后谨慎地给出符合要求的答案。

# 题目材料说明
下面将给你一段网页或者pdf数据解析后的文本，它是可能带有少量html标签的文章小段，并在网页中以零碎的语句呈现。
出题者将为你按顺序标记这些零碎语句，以方便你对文章小段有很好的语义理解。
[注意]出题者还会显式地帮助你分割出文章的`上段`和`下段`，这与你的任务有关。

# 题目任务 —— 找到过渡句
作者为了语义的连贯性，有时候会在上段内容的末尾处补充过渡句。因此，你的任务是判断下面文本的上段部分是否存在过渡句。
[注意]虽然同为过渡句，但是我们定义的过渡句与常规的过渡句不同，我们对过渡句有更加严格的定义，因此请认真理解`题目要求`。

# 题目要求
如`任务`所说，我们对过渡句的要求比较严格，除了对可能的过渡句范围做了划分，还对过渡句有自己的定义：

### 过渡句可选答案(回答选项)
你仅可以在ABCD里选择一个作为过渡句识别的答案：
A. 上段的最后1句是过渡句
B. 上段的最后2句在首尾拼接后是过渡句
C. 上段的最后3句在首尾拼接后是过渡句
D. 无过渡句
[注意]我们划分最后n句为过渡句，不是因为有连续的n句都是过渡句，而是因为这n句共同构成了一个完整的过渡句。

### 过渡句定义
如果从承上启下的角度理解过渡句，那么一篇文章会有海量的过渡句。因此我们收缩了对过渡句的定义，该定义有两部分，一部分是必须遵守的，另一部分是4种过渡句种类：
** 必须遵守的定义 **
1. 该句必须是一个完整句子。
2. 不能有增量信息：增量信息是指一个句子为读者提供了其他句子都没有提供过的信息，例如该句提供了额外的事实描述，例如有其他句子都没有提到的数据、信息、看法。相反地，如果该句是对上文的总结，或者是对下文即将介绍的内容的概述，则不具有增量信息。
** 4种过渡句 **
1. 过渡句式过渡: 有明显的过渡语义，例如`下面我将...`、`接下来...`、`以下...`、`如下...`、`我们请到了3位嘉宾，请他们为我们...`、`让我们来看一看...`、`虽然...，但是...?`、`...,这次真的...吗？`、`应该如何...?`、`...会怎样？`、`我们整理了...分享给大家.`等等。
2. 转折句式过渡: 这种过渡句通过抛出一个转折疑问句｜转折陈述句，然后在下段开始对该转折进行讨论。
[正例]z=`这条路上，商业的部分被大大前置，但仍旧不妨碍猫助和多抓鱼给文艺青年们讲了一个充满理想主义和情怀的故事。`，然后下段介绍猫助和多抓鱼的具体内容。
[正例]z=`xxxxxx, 然而xxx不这么看`，或者z=`尽管如此，我们仍可发现一些规律`，然后下段展开细说。
3. 疑问句式过渡: 抛出一个问题，该问题会引出下段要讲述的对象或者话题。
[正例]z=`新晋AI搜索“黑马”秘塔究竟是啥？`，下段围绕`秘塔`开始介绍，此时它是过渡句。
[负例]z=“伴随着微信小程序的推出，新的流量红利是否可能重新带来O2O的爆发？”，下段并没有对这个问题的展开讨论，因此它不是过渡句。
4. 陈述句式过渡: 满足（1）或（2）任意一个的陈述句就是过渡句：（1）概括了下文将讲述的内容;（2）引出了下文要讲述的对象或事件。
[正例]z=“眼下瓶装水市场再次悄无声息的进入了新一轮商战。”，下段介绍了市场进行商战的情况，因此是过渡句。
[负例]z=“除了解决上面一系列问题之外，数字经济还是我们跟发达国家弯道超车的一个机会点。”，该句可能是为下文讲述“数字经济”做过渡，然而下段没有讲数字经济，而是对全文做了总结，虽然总结内容包括“数字经济”，但是显然不构成过渡。

### 过渡句可能呈现的形式
过渡句在上段的末尾可能被切成多段，因此可能有如下形式:
** 形式1 **
16:那么2023年国庆作为近年来题材最多元的一届，内容与市场都有哪些看点？
17:谁最有可能成为最后赢家？
在这个例子中，16和17是上段的最后2句，17自己可以作为疑问句式过渡，16也是疑问句式过渡，16和17虽然看似是两句话，但实际上组合后是一个大的疑问式过渡句，且没有增量信息，因此最终答案是B。
** 形式2 **
157:以下是PayPal成功的五个主要关键转型
158:。
在这个例子中，157是一个过渡句，但是它的标点被分隔到了158，在157和158拼接后成为一个完整的过渡句，因此答案选B。

### 不是过渡句的例子
虽然有时候一个编号句是过渡句，但是也会有该编号并非一个合法的句子的情况。下面列举几个看似在过渡，实际上不是句子的例子：
** 例1 **
21:（五）项目规模：
该句只是一个小标题，而非一个句子，因此不是过渡句。
** 例2 **
54:9㎡。
该句并不能表达任何意思，甚至无法成为一个句子，除非它与上文拼接后能形成完整的过渡句，否则该句绝对不是过渡句或过渡句的一部分。

### 详细的分析步骤
我们用xyz代表上段的最后三句。请在充分学习“过渡句可能呈现的形式”和“不是过渡句的例子”后，按照step 1～11的分析步骤，一步一步地进行分析。
``` 分析步骤
step 1: 请用比较全面的语言总结下段说了什么, 然后进入step2.
step 2: 分析z是否满足必须遵守的定义？
step 3: 结合0对下段的总结，请分析z是否符合4种过渡句其中的一种？
step 4: 如果1和2答案都是“yes”，则停止之后的分析，且最终答案是A。
step 5: 分析y+z是否满足必须遵守的定义？
step 6: 结合0对下段的总结，请分析y+z是否符合4种过渡句其中的一种？
step 7: 如果4和5答案都是“yes”，则停止之后的分析，且最终答案是B。
step 8: 分析x+y+z是否满足必须遵守的定义？
step 9: 结合0对下段的总结，请分析x+y+z是否符合4种过渡句其中的一种？
step 10: 如果4和5答案都是“yes”，则停止之后的分析，且最终答案是C。
step 11: 最终答案是D。
```

## output format
```json
{"analyse":慢慢思考的分析过程, "answer":最终答案}
"""

def is_BOS(lb, plb):
    if lb == 'BOS' and 'title' in plb:
        return True 
    else:
        return False
    
def extra_trans(s):
    for k,v in html_tag_to_extra_id.items():
        s = s.replace(v, f'<{k}>')
    return s
    
    
    
def create_gpt_messages(sentences_left, 
                        sentences_right, 
                        sentence_numbers_left,
                        sentence_numbers_right):
    system = SYSTEM 
    sentences_left = [extra_trans(s) for s in sentences_left]
    sentences_right = [extra_trans(s) for s in sentences_right]
    candidate_trans = sentences_left[-3:]
    while len(candidate_trans) < 3:
        candidate_trans = [None] + candidate_trans 
    user = {
        'x':candidate_trans[0],
        'y':candidate_trans[1],
        'z':candidate_trans[2],
        "上段":{sn:s for s, sn in zip(sentences_left,sentence_numbers_left)},
        "下段":{sn:s for s, sn in zip(sentences_right,sentence_numbers_right)},
        # "过渡句范文":'\n'.join([f'{index+1} - {st.strip()}' for index, st in enumerate(sample_trans)])
    }
    user = json.dumps(user, ensure_ascii=False)
    return [{'role':'system','content':system},{'role':'user','content':user}]

def gen_trans_gpt_messages_human(line):
    messages = []
    sentences, sentence_numbers, labels = line['sentences'], line['sentence_numbers'], line['labels']
    trans_nums = line.get('transition_sentence_numbers',[])
    packed_trans_nums = []
    tmp_trans = []
    skip = 0
    for t in trans_nums:
        if not tmp_trans or t == tmp_trans[-1]+1:
            tmp_trans.append(t)
        else:
            packed_trans_nums.append(tmp_trans)
            tmp_trans = [t]
    if tmp_trans: packed_trans_nums.append(tmp_trans)
    packed_trans_nums = {i[-1]:i for i in packed_trans_nums}
    
    for index, (s, sn, lb) in enumerate(zip(sentences, sentence_numbers, labels)):
        if skip:
            skip -= 1 
            continue
        if sn in packed_trans_nums.keys() or lb=='BOS':
            if lb=='BOS':
                gt_trans = [sentence_numbers[index-1]] if sentence_numbers[index-1] in trans_nums else []
                index -= 1
            else:
                gt_trans = packed_trans_nums[sn]
            
            left_s = sentences[max(0,index-15): index+1]
            right_s = sentences[index+1: min(len(sentences),index+15)]
            left_sn = sentence_numbers[max(0,index-15): index+1]
            right_sn = sentence_numbers[index+1: min(len(sentences),index+15)]
            message = create_gpt_messages(left_s, right_s, left_sn, right_sn)
            messages.append({
                'messages':message,
                'pre':(left_s, left_sn),
                'suf':(right_s, right_sn),
                'gt_trans':gt_trans,
                'file':line['file']
            })
            skip += 8
    return messages


def gen_trans_gpt_messages(line, 
                        #    sample_trans
                           ):
    messages = []
    sentences, sentence_numbers, labels = line['sentences'], line['sentence_numbers'], line['labels']
    trans_nums = list(set(line.get('transition_sentence_numbers',[])))
    
    pre_BOS = False
    skip = 0
    for idx, lb in enumerate(labels):
        if skip:
            skip -= 1
            continue

        if lb == 'BOS' and not pre_BOS:
            pre_BOS = True
            left = 15 # random.randint(10, 15)
            right = 15 # random.randint(5,8)
            left = max(0, idx-left)
            right = min(len(sentence_numbers), idx + right)
            
            sub_sentences_left = sentences[left:idx]
            sub_sentence_numbers_left = sentence_numbers[left:idx]
            sub_sentences_right = sentences[idx:right]
            sub_sentence_numbers_right = sentence_numbers[idx:right]
            
            _trans_num = [n for n in sub_sentence_numbers_left+sub_sentence_numbers_right if n in trans_nums]

            gpt_messages = create_gpt_messages(sub_sentences_left, 
                                               sub_sentences_right,
                                               sub_sentence_numbers_left, 
                                               sub_sentence_numbers_right,
                                            #    sample_trans
                                               )
            messages.append({'messages':gpt_messages,
                             "pre":(sub_sentences_left,sub_sentence_numbers_left),
                             "suf":(sub_sentences_right, sub_sentence_numbers_right),
                             'gt_trans':_trans_num})
            skip += right//2
        else:
            pre_BOS = False
               
    return messages


if __name__ == '__main__':
    
    
    paths = [
        ('/mnt/user/linzhixin/sentence_classification/data/tar/labeled-pdf-1202/in_level/1.jsonl',
        '/mnt/user/linzhixin/sentence_classification/data/tar/labeled-pdf-1202/in_level/2.jsonl'),
        ('/mnt/user/linzhixin/sentence_classification/data/tar/labeled-web-1202/in_level/1.jsonl',
        '/mnt/user/linzhixin/sentence_classification/data/tar/labeled-web-1202/in_level/2.jsonl')
    ]
    
    # trans_few_shot_store = []
    # for path in paths:
    #     with open(path, 'r', encoding='utf-8') as f:
    #         for line in f.readlines():
    #             line = json.loads(line)
    #             trans_list = line['transition_sentence_numbers']
    #             for s, sn in zip(line['sentences'], line['sentence_numbers']):
    #                 if sn in trans_list and len(s)>5:
    #                     trans_few_shot_store.append(s)
    
    save_dirs = ['/mnt/public/linzhixin/data/sc-filter/1203/pdf-human',
                '/mnt/public/linzhixin/data/sc-filter/1203/web-human',
                ]
    [os.makedirs(dir, exist_ok=True) for dir in save_dirs]
    # 把标注数据也推理一下
    for (p1,p2), sd in zip(paths, save_dirs):
        name = os.path.basename(p1)
        datas = {}
        with open(p1, 'r', encoding='utf-8') as f:
            for line in tqdm(f.readlines()):
                line = json.loads(line)
                datas[line['file']] = line
                
        with open(p2, 'r', encoding='utf-8') as f:
            for line in tqdm(f.readlines()):
                line = json.loads(line)
                file = line['file']
                if file in datas.keys():
                    datas[file]['transition_sentence_numbers'].extend(line['transition_sentence_numbers'])
                    
        # spark
        import random, sys, os, json, math, time

        import findspark

        #指定spark_home为刚才的解压路径,指定python路径
        spark_home = "/mnt/user/linzhixin/mapreduce/spark"
        python_path = "/mnt/user/linzhixin/mapreduce/.venv/bin/python"
        findspark.init(spark_home,python_path)

        import pyspark 
        from pyspark import SparkContext, SparkConf
        conf = SparkConf().setAppName("rdd_tutorial").setMaster("local[16]")
        sc = SparkContext(conf=conf)

        
        for line in datas.values():   
            res = gen_trans_gpt_messages_human(line)
            
            
        with open(os.path.join(sd,"trans_"+name), 'w+') as sf, open(os.path.join(sd, "no-trans_"+name), 'w+') as nsf:  
                for item in res:
                    if item['gt_trans']:
                        json.dump(item, sf, ensure_ascii=False)
                        sf.write('\n')
                    else:
                        json.dump(item, nsf, ensure_ascii=False)
                        nsf.write('\n')

    # assert False 
    
    dirs = ['/mnt/public/deeplang_prod/edu/data-raw/web',
           '/mnt/public/deeplang_prod/edu/data-raw/pdf-data']
    save_dirs = ['/mnt/public/linzhixin/data/sc-filter/1203/web',
                '/mnt/public/linzhixin/data/sc-filter/1203/pdf']
    
    try:
        [os.makedirs(dir) for dir in save_dirs if not os.path.exists(dir)]
    except:
        pass
    
    for dir, save_dir in zip(dirs, save_dirs):
        paths = [os.path.join(dir, name) for name in os.listdir(dir) if '1-part' in name]
        names = [name for name in os.listdir(dir) if '1-part' in name]
        datas = []
        
        for path, name in zip(paths, names):
            print(name)
            with open(path, 'r', encoding='utf-8') as rf, open(os.path.join(save_dir, name),'w') as sf:
                count = 0
                save_num = 0
                for l in tqdm(rf.readlines()):
                    line = json.loads(l)
                    segment = ''.join(line['sentences'])
                    # prompt = filter_prompt_gen(line)
                    count += 1
                    
                    filter = False
                    for item in ['个回复', '赞同来自', '本帖后续']:
                        if item in segment:
                            filter = True
                            break
                    if filter:
                        continue
                    # sample_trans = random.sample(trans_few_shot_store, 3)
                    datas = gen_trans_gpt_messages(line)
                    for data in datas:
                        sf.write(json.dumps({**data, 'name':name, 'index':count}, ensure_ascii=False))
                        sf.write('\n')
                        save_num += 1
                print(f"保存了{save_num}个")